#!/usr/bin/env python
# -*- encoding:UTF-8 -*-

from scipy import *
import bisect
import time, os, subprocess

from pylab import figure, show, rand
from matplotlib.patches import Ellipse
from matplotlib.pyplot import *

#~ from guppy import hpy	# for memory profiling

#Algunas constantes
epsilon = 1E-10

RUNNING_ID = (int(time.time()) % 10000)		# Un id de 4 digitos
DESTDIR = "../videos/%04d"%RUNNING_ID
VIDEO_PREFIX = "t%08.2f.png"
NUMBER_OF_ITERATIONS = 50

# For each iteration
CONF_PRINT_TEF = 1		# Future events table
CONF_PRINT_STATE_VECTOR = 1 	# Sytem state vector (position, velocity, change_times)
CONF_PLOT_DISCRETE_T = 0 	# Plots the system at collision times
# Between iterations
CONF_PLOT_CONTINUOUS_T = 1	# Plots the system at constant time rate (better for animation)
CONF_PLOT_TIMESTEP = 0.5

CONF_MAKE_VIDEO = 1
CONF_PLAY_VIDEO = 0
CONF_OPEN_FOLDER = 1

mycolorfunc = matplotlib.pyplot.cm.gist_rainbow	


def load_vectors(filename):
	"""Devuelve (N, pos, vel, radii): 
	N: Numero de atomos
	pos: vector de posiciones
	vel: vector de velocidades
	radii: vector de radios
	masses: atoms mass vector
	"""
	
	# Read lines as a list
	fh = open(filename, "r")
	lines = fh.readlines()
	fh.close()
	# Filtering blank lines 
	lines = filter(lambda x: not x.isspace(), lines)
	print lines

	#Inicialización del vector de átomos
	N = len(lines)
	nrows = N
	pos = zeros((nrows,3),dtype=double)
	vel = zeros((nrows,3),dtype=double)
	radii = zeros((nrows,),dtype=double)
	masses = zeros((nrows,), dtype=double)

	for r in range(nrows):
		row = lines[r].split()
		if len(row) < 8:
			print "Error: line#",r,"has", len(row), "columns"
			exit(1)
		pos[r] = row[0:3]
		vel[r] = row[3:6]
		radii[r] = row[6]
		masses[r] = row[7]

	return (N, pos, vel, radii, masses)





#Inicialización del vector de átomos
(number_of_atoms, atoms_pos, atoms_vel, atoms_radii, atoms_masses) = load_vectors("atoms.txt")

atoms_time = zeros(number_of_atoms)

container_radius = 11
all_atoms_radius = 1

future = []

class event():
	#-------------------------------------------------------------------
	# Métodos __especiales__
	#-------------------------------------------------------------------

	# in C++: constructor
	def __init__(self, t=0, action=None, params=None):
		self.t, self.action, self.params = t, action, params

	# Representación como string
	# in C++: operator <<
	def __str__(self):
		return 'Event id=%08x t=%g: %s(%s)'%(id(self), self.t, self.action.__name__, self.params)

	# Operador de comparación para ordenación por tiempo
	def __cmp__(self,other):
		return sign(self.t - other.t)
	
	#-------------------------------------------------------------------
	# Métodos propios
	#-------------------------------------------------------------------
	def execute(self):
		#~ print "Calling %s(%s)..."%(self.action.__name__, self.params)
		self.action(self.t, self.params)



def estimate_collision_time_with_container(idAtom):
	p = atoms_pos[idAtom] 		# Position vector
	v = atoms_vel[idAtom]		# Velocity vector
	t = atoms_time[idAtom]		# Change time vector
	radius = atoms_radii[idAtom]
	
	# Coeficientes de la ecuación de segundo grado
	# (p.x+v.x*t)^2 +(p.y+v.y*t)^2 +(p.z+v.z*t)^2 = (container_radius - atom_radius)^2
	# t^2 (v.x^2+vy^2+vz^2) + t 2(vx.px + vy.py + vz.pz) + (px^2+py^2+pz^2)-(container_radius<-all_atoms_radius)^2 = 0
	# t^2 (v.v) + t (2.p.v) + (p.p-(container_radius<-all_atoms_radius)^2) = 0
	# (dot = producto punto o interno o euclideo)
	a = dot(v,v)
	b = 2*dot(p,v)
	c = dot(p, p)-(container_radius-radius)**2
	
	det = b**2-4*a*c
	if det > epsilon:
		x1 = (-b + sqrt(det))/(2*a)
		x2 = (-b - sqrt(det))/(2*a)
		
		# Una raiz debe ser positiva y la otra negativa o cero
		#~ assert x1*x2 < epsilon 	# comparación con tolerancia
	
		delta_t = max(x1,x2)
		
		return t + delta_t
	else:
		# Si la velocidad es 0 o si no está dentro del container, 
		# podría no chocar (caso A3)
		return -1				
	

def collide_with_container(t, params):
	print "<<< Colliding with container "  , params

	idAtom = params[0]
	
	p = atoms_pos[idAtom] 		# Position vector
	v = atoms_vel[idAtom]		# Velocity vector
	last_t = atoms_time[idAtom]	# Time of last trajectory change
	
	# Nueva posición y velocidad
	p = p + v*(t-last_t)
	v = v - (2*dot(p,v)/dot(p,p))*p

	atoms_pos[idAtom] = p
	atoms_vel[idAtom] = v
	atoms_time[idAtom] = t
	#~ print '\n\nAt %f, atom #%d \npos=%s\nvel=%s'%(t,idAtom,str(p),str(v))
	#~ !print '\t%8.3f %4d %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f'%(t,idAtom,p[0],p[1],p[2],v[0],v[1],v[2])
	
	# Purgar eventos
	# Limpieza de eventos con id actual
	for e in future:
		if (idAtom in e.params and e.t > t):
			print ">>> Removing ", e
			# future2.remove(e)
	future[:]  = [x for x in future if not( idAtom in x.params and x.t > t)]

	
	
	# Estimar tiempo de colisión con la nueva trayectoria y agendar
	schedule_collision_events(idAtom)
				
def estimate_collision_time_with_atoms(idAtom, r=range(number_of_atoms)):
	# Segun Rapaport04, p392
	# Se aplica la formula 14.2.2
	pi_ = atoms_pos[idAtom] 	# Position vector
	vi = atoms_vel[idAtom]		# Velocity vector
	ti = atoms_time[idAtom]
	
	possible_collisions = []
	
	for idAtom_partner in r:
		
		if idAtom_partner==idAtom:
			continue	
			
		diam = atoms_radii[idAtom]+atoms_radii[idAtom_partner]
			
		pj_ = atoms_pos[idAtom_partner] 		# Position vector
		vj = atoms_vel[idAtom_partner]		# Velocity vector
		tj = atoms_time[idAtom_partner]
		
		tref = max(ti,tj)

		pi = pi_ + vi*(tref-ti)
		pj = pj_ + vj*(tref-tj)
		
		r = pi-pj
		v = vi-vj
		b = dot(r,v)
		det = b*b-dot(v,v)*(dot(r,r)-diam**2)	

		if (b<0 and det>=0): #condicion segun bibliografia necesaria
			delta_t = (-b-sqrt(det))/dot(v,v)
			if (delta_t>=0):
				possible_collisions.append( (tref+delta_t, idAtom_partner) )
	
	return possible_collisions
				
	
		

def collide_between_atoms(dest_t, atoms_pair):
	print "<<< Colliding atoms...", atoms_pair
	id0, id1 = atoms_pair
	pp = ([atoms_pos[id0],atoms_pos[id1]]) 		# Position vector
	vv = ([atoms_vel[id0],atoms_vel[id1]])		# Velocity vector
	tt = (atoms_time[id0],atoms_time[id1])		# Time of last trajectory change
	diam = atoms_radii[id0]+atoms_radii[id1]
	mm = (atoms_masses[id0], atoms_masses[id1]) 
	
	# Nueva posición
	pp[0] += vv[0]*(dest_t - tt[0])
	pp[1] += vv[1]*(dest_t - tt[1])
	#~ print "@@", pp[0], pp[1]
	
	# Nueva velocidad
	r = pp[0]-pp[1]
	v = vv[0]-vv[1]
	b = dot(r,v)
	#~ det=dot(b,b)-dot(v,v)*(dot(r,r)-diam**2)
	#~ print "@@", vv[0], vv[1], v
	det=b**2-dot(v,v)*(dot(r,r)-diam**2)	
	delta_v = (-b/diam**2)*r
	
	vv[0] += delta_v * 2*mm[1]/(mm[0]+mm[1])
	vv[1] -= delta_v * 2*mm[0]/(mm[0]+mm[1])

	#~ atoms_pos[id0] = pp[0]
	#~ atoms_pos[id1] = pp[1]
	#~ atoms_vel[id0],atoms_vel[id1] = vv
	atoms_time[id0],atoms_time[id1] = dest_t, dest_t
	
	
	# Limpieza de eventos con id actual
	# TODO: Considerar epsilon en la comparacion de tiempos
	for e in future:
		if ((id0 in e.params or id1 in e.params) and e.t > dest_t):
			print ">>> Removing ", e
			# future.remove(e)
	future[:]  = [x for x in future if not ((id0 in x.params or id1 in x.params) and x.t > dest_t)] 
        
	
	#~ print '\n\nAt %f, atom #%d \npos=%s\nvel=%s'%(t,idAtom,str(p),str(v))
	for i in (0,1):
		#~ !print '\t%8.3f %4d %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f'%(tt[i],atoms_pair[i],pp[i][0],pp[i][1],pp[i][2],vv[i][0],vv[i][1],vv[i][2])
	
		# Estimar tiempo de colisión con la nueva trayectoria y agendar
		schedule_collision_events(atoms_pair[i])

def schedule_collision_events(idAtom):
	t = estimate_collision_time_with_container(idAtom)
	if t > 0:
                e = event(t, collide_with_container, (idAtom,))
                print "%%% Scheduling ", e
		bisect.insort(future, e)
        
	collisions = estimate_collision_time_with_atoms(idAtom)
	for (t,idPartner) in collisions:
                e = event(t, collide_between_atoms,(idAtom, idPartner) )
                print "%%% Scheduling ", e
                bisect.insort(future, e)
		

def print_header():
	print '#Iter\t  Time   Atom       Position             Velocity      '
	print '=====\t======== ==== ==================== ===================='

def print_header2():
	print '#Iter\t  Time   Atom       Position             Velocity          Energy'
	print '=====\t======== ==== ==================== ==================== ==========='


	
def plot_state(instant_t):
        fig = figure()
        ax = fig.add_subplot(111, aspect='equal')
        ax.set_xlim(-12, 12)
        ax.set_ylim(-12, 12)
        
        container = Ellipse(xy=(0,0), width=22, height=22)
        container.set_alpha(0.1)
        ax.add_artist(container)

	for idAtom in range(number_of_atoms):
		pp = atoms_pos[idAtom] 		# Position vector
		v = atoms_vel[idAtom]		# Velocity vector
		t = atoms_time[idAtom]		# Change time vector
		p = pp +v*(instant_t-t)
		d = 2*atoms_radii[idAtom]	# Diameter
		ax.add_artist(Ellipse(xy=p[0:2], width=d, height=d, color=mycolorfunc(1.0*idAtom/number_of_atoms)))

        fig.savefig(DESTDIR+"/"+VIDEO_PREFIX%instant_t)
        close(fig)


if __name__ == "__main__":
	#~ DESTDIR = "video/%04d"%RUNNING_ID
	os.makedirs(DESTDIR)

	# Se cargan todos los átomos con sus estimaciones iniciales
	for idAtom in range(number_of_atoms):
		t = estimate_collision_time_with_container(idAtom)
		if t >= 0:
			e = event(t, collide_with_container, (idAtom,))
			bisect.insort(future, e)
			#~ future_events.append(e)
        
	for idAtom in range(number_of_atoms):
		collisions = estimate_collision_time_with_atoms(idAtom,range(idAtom+1,number_of_atoms))
		for (t,idPartner) in collisions:
			bisect.insort(future, event(t, collide_between_atoms,(idAtom, idPartner) ))
			

	# Se ejecuta el scheduler
	t = 0	
	plot_state(t)

	for i in range(NUMBER_OF_ITERATIONS):
		if CONF_PRINT_STATE_VECTOR:
			print "========================================"
			print_header2()
			tot_energy = 0.0
			for ii in range(number_of_atoms):
				pp = atoms_pos[ii] 		# Position vector
				vv = atoms_vel[ii]		# Velocity vector
				tt = atoms_time[ii]	# Time of last trajectory change
				mm = atoms_masses[ii]
				energy = mm*dot(vv,vv)
				tot_energy += energy
				print '=\t%8.3f %4d %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f'%(tt,ii,pp[0],pp[1],pp[2],vv[0],vv[1],vv[2],energy)
			print "======================================== Total Energy = ", tot_energy
		
		if CONF_PRINT_TEF:
			print "==== Future events table ========"
			for e in future:
				print "tef>", e
		
                e = future.pop(0)
		print '<*> Event OrderNo #%d, id=%08x, at time=%f'%(i, id(e), e.t)
		
		if CONF_PLOT_CONTINUOUS_T:
			while t < e.t:
				plot_state(t)
				t += CONF_PLOT_TIMESTEP
			
		if CONF_PLOT_DISCRETE_T:
			plot_state(e.t)
			
		e.execute()
		
	
	if CONF_MAKE_VIDEO:
		print 	
		os.chdir(DESTDIR)
		subprocess.Popen(["mencoder", "-o", "%04d.avi"%RUNNING_ID, "mf://*.png", "-ovc", "lavc"])
		if CONF_PLAY_VIDEO:
			subprocess.Popen(["mplayer", "-loop", "0", "%04d.avi"%RUNNING_ID])

		if CONF_OPEN_FOLDER:
			subprocess.Popen(["nautilus", "."])
			
	#~ print "===Memory usage==="
	#~ print hpy().heap() # Show current reachable heap
